import torch
import numpy as np
from torchvision import transforms
from PIL import Image,ImageDraw
from models.CnnNet import CnnNet
def predict(image_path):
    """模型推理测试"""
    #加载模型
    model = CnnNet()
    model.load_state_dict(torch.load('./models/cnnnet.pt', weights_only=True))
    model.eval()

    #预处理图像
    transform = transforms.Compose([
        transforms.Lambda(auto_invert), #由于模型训练时使用的MNIST数据集是黑底白字的图片（背景为0，数字为1）,如果输入图片是白底彩字，导致预处理后的图像与训练数据分布不一致，所以智能背景检测（自动处理黑白背景）
        transforms.Grayscale(),  # 确保转为灰度
        transforms.Resize((28, 28)),  # 强制调整尺寸
        transforms.ToTensor(),
        transforms.Normalize((0.1307,),(0.3081,))
    ])
     #确保为MNIST同格式的灰度图
    image = Image.open(image_path)
    if image.mode != 'L':
            image = image.convert('L')
    # 应用预处理
    image = transform(image).unsqueeze(0) #[1,1,28,28]
    # 验证输入尺寸
    if image.shape != torch.Size([1, 1, 28, 28]):
        raise ValueError(f"输入尺寸异常，得到 {image.shape}，期望 [1,1,28,28]")
    #推理
    with torch.no_grad():
        output = model(image)
        prob = torch.softmax(output,dim=1) #转为概率
        confidence,predicted = torch.max(prob,1)

    #print(f"Predicted: {predicted.item()} (Confidence:{confidence.item():.2%})")
    print("\n推理结果：")
    print(f"预测数字：{predicted.item()}")
    print(f"可信度：{confidence.item()*100:.2f}%")
    
    # 输出所有类别概率（可选）
    print("\n预测概率：")
    for i, p in enumerate(prob.squeeze().tolist()):
        print(f"数字 {i}: {p*100:.2f}%")

def auto_invert(image):
    """自动检测并反转背景"""
    np_img = np.array(image)
    if np_img[0,0] > 128:  # 检测左上角像素亮度
        return Image.fromarray(255 - np_img)
    return image

def sampleImage():
    """生成测试图片"""
    # 1. 创建一个空白数字图像（背景为黑色）
    image = Image.new('L', (28, 28), color=0)  # 'L'表示灰度图
    draw = ImageDraw.Draw(image)

    # 2. 随机绘制一个数字（0-9）
    digit = np.random.randint(0, 10)  # 随机选择数字
    draw.text((10, 5), str(digit), fill=255)  # 白色数字

    # 3. 保存图像
    image.save("test_image.png")
    print(f"已生成测试图像: test_image.png (数字{digit})")

    # 显示图像（可选）
    image.show()

if __name__ == "__main__":
    # sampleImage()
    predict('./data/test/OIP4.jpg')